import xml.dom.minidom as dom
import Image
import os
import os.path
import sys

'''

usage:

    python vocxml.py ~/caffe/data/eagLeNet/plane/all_resized/all_processed

图片：
    0xxxxx.jpg : plane
    1xxxxx.jpg : space

'''

def createNode(doc,nodename,text=''):
    if (text==''):
        node = doc.createElement(nodename)
    else:
        node=doc.createElement(nodename)
        node.appendChild(doc.createTextNode(text))
    return node

def addChildren(root,*nodelist):
    for i in nodelist:
        root.appendChild(i)

def formatNode(doc,root,namelist,textlist):
    if(len(namelist)!=len(textlist)): return
    for i in range(len(namelist)):
        node=createNode(doc,str(namelist[i]),str(textlist[i]))
        root.appendChild(node)

def generateNum(num):
    num=str(num)
    length=len(num)
    name='1'+((5-length)*'0')+num
    return name

def main(path):
    num=0
    savepath = os.path.join(path,'Annotation')

    if os.path.exists(savepath)==False :
        os.mkdir(savepath)
    for parent,dirnames,filenames in os.walk(path):
        for filename in filenames:
            print path+filename

            if(filename.endswith('jpg')):
                num += 1
                im=Image.open(path+'/'+filename).convert('RGB')
                x,y=im.size

                count=filename[:-4]
                name=filename
                root=generateXML(name,x,y)

                with open(savepath+'/'+count+'.xml','w') as f:
                    f.write(root.toprettyxml())



def generateXML(name,x,y):

    doc=dom.Document()
    root=createNode(doc=doc,nodename='annotation')
    doc.appendChild(root)

    folder=createNode(doc=doc,nodename='folder',text='eag')
    filename=createNode(doc,'filename',name)

    source=createNode(doc,'source')
    sourcenames=['database',
                 'annotation',
                 'image',
                 'flickrid']
    sourcetexts=['the eag Database',
                 'eag',
                 'flickr',
                 'NULL']
    formatNode(doc,source,sourcenames,sourcetexts)

    owner=createNode(doc,'owner')
    ownernames=['flickrid',
                'name']
    ownertexts=['daggerage',
                'plane']
    formatNode(doc,owner,ownernames,ownertexts)

    size=createNode(doc,'size')
    sizenames=['width',
               'height',
               'depth']
    sizetexts=[x,y,3]
    formatNode(doc,size,sizenames,sizetexts)

    segmented=createNode(doc,'segmented','0')

    if(name.startswith('0')):
        objname='plane'
    else:
        objname='space'
    object=createNode(doc,'object')
    objectnames=['name',
                 'pose',
                 'truncated',
                 'difficult']
    objecttext=[objname,
                'Unspecified',
                1,
                0]

    bndbox=createNode(doc,'bndbox')
    bndboxnames=['xmin',
                 'ymin',
                 'xmax',
                 'ymax',]
    bndboxtexts=[0,0,x,y]

    formatNode(doc, object, objectnames, objecttext)
    formatNode(doc,bndbox,bndboxnames,bndboxtexts)
    object.appendChild(bndbox)

    addChildren(root,folder,filename,source,owner,size,segmented,object)

    return root

if __name__ == '__main__':
    main(sys.argv[1])
